﻿// Copyright  (c) 2011, The Outercurve Foundation. 
// The software is licensed under the New BSD License (the "License"); 
// you may not use the software except in compliance with the License. 

namespace NMTopProtocols
{
    using System;
    using System.Collections.Generic;
    using System.Linq;
    using System.Text;

    /// <summary>
    /// Used to create an internal representation of how protocols are related and the count of that type of protocol we've seen
    /// in a trace
    /// </summary>
    public class ProtocolTree
    {
        /// <summary>
        /// Used to give each instance of the class a unique ID
        /// </summary>
        private static uint internalIDCounter = 0;

        /// <summary>
        /// This tree's unique ID
        /// </summary>
        private uint myid;

        /// <summary>
        /// Name of the protocol leaf
        /// </summary>
        private string name;

        /// <summary>
        /// Number of times we've seen this protocol
        /// </summary>
        private uint count;

        /// <summary>
        /// Number of bytes of this protocol we've seen
        /// </summary>
        private ulong bytes;

        /// <summary>
        /// Number of bytes through reassembly we've seen
        /// </summary>
        private ulong reassembledBytes;

        /// <summary>
        /// Number of times we've seen this protocol reassembled
        /// </summary>
        private uint reassembledCount;

        /// <summary>
        /// Keep track of this node's children
        /// </summary>
        private Dictionary<string, ProtocolTree> branches = new Dictionary<string, ProtocolTree>();

        /// <summary>
        /// Keep track of this node's parent
        /// </summary>
        private ProtocolTree parent;

        /// <summary>
        /// The Protocol Path of this node
        /// </summary>
        private string path;

        /// <summary>
        /// Initializes a new instance of the ProtocolTree class.
        /// This is a root node for a tree.
        /// </summary>
        /// <param name="name">Name of Root Protocol</param>
        public ProtocolTree(string name) : this(name, null)
        {
        }

        /// <summary>
        /// Initializes a new instance of the ProtocolTree class.
        /// This is a child of an existing tree.
        /// </summary>
        /// <param name="name">Name of Protocol</param>
        /// <param name="parent">Parent Protocol Node</param>
        public ProtocolTree(string name, ProtocolTree parent)
        {
            this.parent = parent;
            this.name = name;
            this.myid = ProtocolTree.internalIDCounter;
            ProtocolTree.internalIDCounter++;

            // Calculate Path String
            string path = this.name;
            ProtocolTree p = this.parent;
            while (p != null)
            {
                path = p.name + "." + path;

                p = p.parent;
            }

            this.path = path;
        }

        /// <summary>
        /// Gets the Name of this Protocol
        /// </summary>
        public string ProtocolName
        {
            get { return this.name; }
        }

        /// <summary>
        /// Gets the quantified path of this Protocol
        /// </summary>
        public string Path
        {
            get { return this.path; }
        }

        /// <summary>
        /// Gets the count of frames of this Protocol
        /// </summary>
        public uint Frames
        {
            get { return this.count; }
        }

        /// <summary>
        /// Gets the count of reassembled frames of this Protocol
        /// </summary>
        public uint ReassembledCount
        {
            get { return this.reassembledCount; }
        }

        /// <summary>
        /// Gets the number of bytes of this Protocol
        /// </summary>
        public ulong Bytes
        {
            get { return this.bytes; }
        }

        /// <summary>
        /// Gets the number of reassembled bytes of this Protocol
        /// </summary>
        public ulong ReassembledBytes
        {
            get { return this.reassembledBytes; }
        }

        /// <summary>
        /// Gets the ID of this tree node
        /// </summary>
        public uint ID
        {
            get { return this.myid; }
        }

        /// <summary>
        /// Gets the children of this node
        /// </summary>
        public ProtocolTree[] Children
        {
            get { return this.branches.Values.ToArray(); }
        }
        
        /// <summary>
        /// Adds a child to the tree constructing any nodes needed inbetween
        /// </summary>
        /// <param name="protocolpath">Fully Quantified string of Protocol Path separated by periods</param>
        /// <param name="length">Number of bytes to add</param>
        /// <returns>The newly constructed leaf</returns>
        public ProtocolTree AddProtocol(string protocolpath, uint length)
        {
            if (protocolpath.Contains("."))
            {
                string n = protocolpath.Split('.')[0];

                if (!this.branches.ContainsKey(n))
                {
                    this.branches[n] = new ProtocolTree(n, this);
                }

                return this.branches[n].AddProtocol(protocolpath.Substring(n.Length + 1), length);
            }
            else
            {
                if (!this.branches.ContainsKey(protocolpath))
                {
                    this.branches[protocolpath] = new ProtocolTree(protocolpath, this);
                }

                this.branches[protocolpath].AddFrame(length);

                return this.branches[protocolpath];
            }
        }

        /// <summary>
        /// Adds a frame of particular length to this node
        /// </summary>
        /// <param name="length">Number of bytes to add</param>
        public void AddFrame(uint length)
        {
            this.IncrementCount();
            this.IncrementBytes(length);
        }

        /// <summary>
        /// Adds a reassembled frame to this node (also increments regular count)
        /// </summary>
        /// <param name="length">Number of bytes to add</param>
        public void AddReassembledFrame(uint length)
        {
            this.AddFrame(length);
            this.reassembledBytes += length;
            this.reassembledCount += 1;
        }

        /// <summary>
        /// Calculates the sum of bytes of this node and all its children recursively
        /// </summary>
        /// <returns>Total number of bytes</returns>
        public ulong GetBytesWithAllChildren()
        {
            ulong b = this.bytes;

            foreach (ProtocolTree node in this.branches.Values)
            {
                b += node.GetBytesWithAllChildren();
            }

            return b;
        }

        /// <summary>
        /// Calculates the sum of the count of this node and all its children recursively
        /// </summary>
        /// <returns>Total number of frames</returns>
        public ulong GetFramesWithAllChildren()
        {
            ulong b = this.count;

            foreach (ProtocolTree node in this.branches.Values)
            {
                b += node.GetFramesWithAllChildren();
            }

            return b;
        }

        /// <summary>
        /// Searches this node and its Children for the specified path and returns the number of bytes in that node
        /// </summary>
        /// <param name="searchpath">Fully Quantified Path</param>
        /// <returns>Number of bytes of the child</returns>
        public ulong GetChildBytes(string searchpath)
        {
            if (this.path == searchpath)
            {
                return this.bytes;
            }

            foreach (ProtocolTree node in this.branches.Values)
            {
                ulong r = node.GetChildBytes(searchpath);

                if (r != 0)
                {
                    return r;
                }
            }

            return 0;
        }

        /// <summary>
        /// Finds the child with the largest number of bytes
        /// </summary>
        /// <returns>Child node</returns>
        public ProtocolTree GetLargestChildBytes()
        {
            ProtocolTree cur = this;

            foreach (ProtocolTree node in this.branches.Values)
            {
                ProtocolTree n = node.GetLargestChildBytes();

                if (n.bytes > cur.bytes)
                {
                    cur = n;
                }
            }

            return cur;
        }

        /// <summary>
        /// Returns a list of all the leaves from this node
        /// </summary>
        /// <returns>list of leaves</returns>
        public List<ProtocolTree> GetLeaves()
        {
            List<ProtocolTree> cur = new List<ProtocolTree>();

            if (this.branches.Count > 0)
            {
                foreach (ProtocolTree node in this.branches.Values)
                {
                    cur.AddRange(node.GetLeaves());
                }
            }
            else
            {
                cur.Add(this);
            }

            return cur;
        }

        /// <summary>
        /// Returns a flat list of all nodes of the tree
        /// </summary>
        /// <returns>List of nodes</returns>
        public List<ProtocolTree> GetFlatTree()
        {
            List<ProtocolTree> cur = new List<ProtocolTree>();

            if (this.branches.Count > 0)
            {
                foreach (ProtocolTree node in this.branches.Values)
                {
                    cur.AddRange(node.GetFlatTree());
                }
            }

            cur.Add(this);

            return cur;
        }

        /// <summary>
        /// Increments the frame count of this protocol
        /// </summary>
        private void IncrementCount()
        {
            this.count++;
        }

        /// <summary>
        /// Increments the number of bytes of this protocol
        /// </summary>
        /// <param name="length">Number of bytes to add</param>
        private void IncrementBytes(uint length)
        {
            this.bytes += length;
        }
    }
}
